from datetime import datetime
from motor.motor_asyncio import AsyncIOMotorDatabase
from models.rbs import RB
from .base import BaseRepository

class RBSRepository(BaseRepository[RB]):  # 直接使用RBS模型
    def __init__(self, db: AsyncIOMotorDatabase):
        super().__init__("rbs", db)

    async def get_by_rb_id(self, rb_id: str) -> RB:
        """
        根据rbid获取数据
        """
        return await self.collection.find_one({"rbId": rb_id})

    async def create_or_update(self, rb: RB) -> RB:
        """
        创建或更新数据
        """
        query = {"rbId": rb.rbId}
        rbs_dict = rb.model_dump()
        
        # 准备更新数据
        new_data = {k: v for k, v in rbs_dict.items() if k != "createTime"}
        new_data["updateTime"] = datetime.now()
        
        # 当插入新文档时设置createtime
        insert_data = {"createTime": rb.createTime}
        
        # 使用 update_one 方法，upsert=True 表示如果不存在则插入，存在则更新
        await self.collection.update_one(
            query,
            {
                "$set": new_data,
                "$setOnInsert": insert_data
            },
            upsert=True
        )
        
        # 返回更新后的文档
        return await self.get_by_rb_id(rb.rbId)    